using Microsoft.Win32;
using System;
using System.Diagnostics;
using System.Security.Principal;
using System.Windows.Forms;

namespace demo
{
    /// <summary>
    /// USB存储设备控制器
    /// 提供USB存储设备的启用/禁用功能和管理员权限检测
    /// </summary>
    public class UsbStorageController
    {
        #region 常量定义
        private const string UsbStorageRegistryKey = @"SYSTEM\CurrentControlSet\Services\USBSTOR";
        private const string StartValueName = "Start";
        private const int USB_DISABLED = 4;  // 禁用状态
        private const int USB_MANUAL = 3;    // 手动启动状态
        #endregion

        #region USB存储设备控制功能
        
        /// <summary>
        /// 禁用USB存储设备
        /// </summary>
        /// <returns>操作结果</returns>
        public UsbOperationResult DisableUsbStorageDevices()
        {
            try
            {
                using (RegistryKey key = Registry.LocalMachine.OpenSubKey(UsbStorageRegistryKey, true))
                {
                    if (key != null)
                    {
                        // 将 USBSTOR 服务的启动类型设置为 4（禁用）
                        key.SetValue(StartValueName, USB_DISABLED, RegistryValueKind.DWord);
                        return new UsbOperationResult
                        {
                            Success = true,
                            Message = "已成功禁用 USB 移动存储设备。此设置将在重启后完全生效。"
                        };
                    }
                    else
                    {
                        return new UsbOperationResult
                        {
                            Success = false,
                            Message = "无法找到 USBSTOR 注册表项。请确认您的系统支持此功能。"
                        };
                    }
                }
            }
            catch (Exception ex)
            {
                return new UsbOperationResult
                {
                    Success = false,
                    Message = $"禁用 USB 移动存储设备时发生错误: {ex.Message}"
                };
            }
        }

        /// <summary>
        /// 启用USB存储设备
        /// </summary>
        /// <returns>操作结果</returns>
        public UsbOperationResult EnableUsbStorageDevices()
        {
            try
            {
                using (RegistryKey key = Registry.LocalMachine.OpenSubKey(UsbStorageRegistryKey, true))
                {
                    if (key != null)
                    {
                        // 将 USBSTOR 服务的启动类型设置为 3（手动）
                        key.SetValue(StartValueName, USB_MANUAL, RegistryValueKind.DWord);
                        return new UsbOperationResult
                        {
                            Success = true,
                            Message = "已成功启用 USB 移动存储设备。此设置将在重启后完全生效。"
                        };
                    }
                    else
                    {
                        return new UsbOperationResult
                        {
                            Success = false,
                            Message = "无法找到 USBSTOR 注册表项。请确认您的系统支持此功能。"
                        };
                    }
                }
            }
            catch (Exception ex)
            {
                return new UsbOperationResult
                {
                    Success = false,
                    Message = $"启用 USB 移动存储设备时发生错误: {ex.Message}"
                };
            }
        }

        /// <summary>
        /// 获取当前USB存储设备状态
        /// </summary>
        /// <returns>USB设备状态</returns>
        public UsbDeviceStatus GetUsbStorageDeviceStatus()
        {
            try
            {
                using (RegistryKey key = Registry.LocalMachine.OpenSubKey(UsbStorageRegistryKey, false))
                {
                    if (key != null)
                    {
                        object value = key.GetValue(StartValueName);
                        if (value != null)
                        {
                            int startValue = Convert.ToInt32(value);
                            return startValue == USB_DISABLED ? UsbDeviceStatus.Disabled : UsbDeviceStatus.Enabled;
                        }
                    }
                }
            }
            catch (Exception)
            {
                // 发生异常时返回未知状态
            }
            return UsbDeviceStatus.Unknown;
        }

        #endregion

        #region 管理员权限检测
        
        /// <summary>
        /// 检查当前程序是否以管理员权限运行
        /// </summary>
        /// <returns>是否具有管理员权限</returns>
        public bool IsAdministrator()
        {
            try
            {
                var identity = WindowsIdentity.GetCurrent();
                var principal = new WindowsPrincipal(identity);
                return principal.IsInRole(WindowsBuiltInRole.Administrator);
            }
            catch (Exception)
            {
                return false;
            }
        }

        /// <summary>
        /// 以管理员权限重启当前应用程序
        /// </summary>
        /// <returns>重启操作结果</returns>
        public UsbOperationResult RestartAsAdministrator()
        {
            var startInfo = new ProcessStartInfo
            {
                UseShellExecute = true,
                WorkingDirectory = Environment.CurrentDirectory,
                FileName = Application.ExecutablePath,
                Verb = "runas"
            };

            try
            {
                Process.Start(startInfo);
                Application.Exit();
                return new UsbOperationResult
                {
                    Success = true,
                    Message = "正在以管理员权限重启应用程序..."
                };
            }
            catch (Exception ex)
            {
                return new UsbOperationResult
                {
                    Success = false,
                    Message = $"请求管理员权限失败: {ex.Message}"
                };
            }
        }

        #endregion
    }

    #region 辅助类和枚举
    
    /// <summary>
    /// USB操作结果
    /// </summary>
    public class UsbOperationResult
    {
        /// <summary>
        /// 操作是否成功
        /// </summary>
        public bool Success { get; set; }
        
        /// <summary>
        /// 操作结果消息
        /// </summary>
        public string Message { get; set; }
    }

    /// <summary>
    /// USB设备状态枚举
    /// </summary>
    public enum UsbDeviceStatus
    {
        /// <summary>
        /// 已启用
        /// </summary>
        Enabled,
        
        /// <summary>
        /// 已禁用
        /// </summary>
        Disabled,
        
        /// <summary>
        /// 未知状态
        /// </summary>
        Unknown
    }
    
    #endregion
}